#include "copycat.h"
static CUProgress progress;


//////////////////////////////////////////////////////////////////////////////
bool CCopycatTrain::Train( char *parse_fn, char *giza_fn, char *align_fn, char *out_fn )
{
	if( LoadCorpus( parse_fn, giza_fn, align_fn ) == false )	return false;
	if( out_files.Prepare( out_fn ) == false ) 	return false;

	if( ExtractSrcSubtree() == false )	return false;

//	if( ExtractRuleCross() == false )	return false;

	if( ExtractRuleRF() == false )	return false;

	return true;
}


//////////////////////////////////////////////////////////////////////////////
bool CCopycatTrain::ExtractSrcSubtree()
{
	stree.Clear();

	for( int i=0; i<num_sent; i++ )	{
		fprintf( stderr, "\r[Extract Src Subtree]   %d / %d sentences processed (# Subtree: %d)", i, num_sent, stree.GetSize() );

		CSent &s = sent[i];


		// load parse tree
		CTree t ;
		t.Init();
		if( t.LoadTree( s.parse ) == false )	return false;
		if( SetNodeSpan( s, t ) == false )	{
			fprintf( stderr, "error %d\n", i );
			return false;
		}

		// for each node
		for( int j=0; j<t.num_node; j++ )	{
			CNode &n = t.node[j];

			if( MCH == "2" )	{ if( n.num_child != 2 )	continue; }
			else if( MCH == "23" )	 { if( n.num_child == 2  || n.num_child ==3 )	{ } else	{ continue; } }
			else { if( n.num_child <= 1 ) continue;  }


			if( IsAllChildrenAligned( t, n ) == false )	continue;

			// extract src patterns
			int num_cross = GetNumCross( t, n, s );
			CUString rule_str = GetRuleStr( t, n );
			stree.Add( rule_str, num_cross );
		}

	}
	fprintf( stderr, "\r[Extract Src Subtree]   %d / %d sentences processed (# Subtree: %d)\n",
					num_sent, num_sent, stree.GetSize() );

	if( stree.SaveToFile( out_files.file_src_info.GetFP() ) == false )	{
		fprintf( stderr, "[Fail to Save]\n" );
		return false;	
	}

	return true;	
}


//////////////////////////////////////////////////////////////////////////////
void CCopycatTrain::RecoverReorder( CNode &r_node, CNode &org_node )
{
	for( int i=0; i<r_node.num_child; i++ )	{
		r_node.child[i] = org_node.child[i];
	}

}


//////////////////////////////////////////////////////////////////////////////
bool CCopycatTrain::ExtractRuleCross()
{
	for( int i=0; i<num_sent; i++ )	{
		fprintf( stderr, "\r[Extract Rule: CR   ]   %d / %d sentences processed (# Rules: %d)", i, num_sent, rule_cross.GetNumRules() );

		CSent &s = sent[i];

		// load tree
		CTree t, t2;
		t.Init(); t2.Init();
		if( t.LoadTree( s.parse ) == false )	return false;
		if( t2.LoadTree( s.parse ) == false )	return false;
		SetNodeSpan( s, t );
		SetNodeSpan( s, t2 );

		// store src statics
		for( int j=0; j<t.num_node; j++ )	{
			CNode &n = t.node[j];
			CNode &n2 = t2.node[j];

			//if( n.num_child != 2 )	continue;
			if( n.num_child == 2 )	{
//				if( IsAllChildrenAligned( t, n ) == false )	continue;

				CRuleCross new_rule;
				new_rule.before = GetRuleStr( t, n );
				new_rule.before_cross = GetNumCross( t2, n2,s );
				new_rule.freq = 1;

				// reorder
				int tmp = n2.child[0];
				n2.child[0] = n2.child[1];
				n2.child[1] = tmp;

				// insert
				new_rule.after = "A1 A0";
				new_rule.after_cross = GetNumCross( t2, n2, s );
				rule_cross.Add( new_rule );
				RecoverReorder( n2, n );
			}
			else if( n.num_child == 3 )	{
				continue;
//				if( IsAllChildrenAligned( t, n ) == false )	continue;

				CRuleCross new_rule;
				new_rule.before = GetRuleStr( t, n );
				new_rule.before_cross = GetNumCross( t2, n2,s );
				new_rule.freq = 1;

				int tmp;

				// reorder type1: 0 1 2 -> 0 2 1
				tmp = n2.child[1]; n2.child[1] = n2.child[2]; n2.child[2] = tmp;
				new_rule.after = "A0 A2 A1";
				new_rule.after_cross = GetNumCross( t2, n2, s );
				rule_cross.Add( new_rule );
				RecoverReorder( n2, n );

				// reorder type2: 0 1 2  -> 1 0 2
				tmp = n2.child[0]; n2.child[0] = n2.child[1]; n2.child[1] = tmp;
				new_rule.after = "A1 A0 A2";
				new_rule.after_cross = GetNumCross( t2, n2, s );
				rule_cross.Add( new_rule );
				RecoverReorder( n2, n );

				// reorder type3: 0 1 2  -> 2 0 1
				tmp = n2.child[2]; n2.child[2] = n2.child[1]; n2.child[1] = n2.child[0]; n2.child[0] = tmp;
				new_rule.after = "A2 A0 A1";
				new_rule.after_cross = GetNumCross( t2, n2, s );
				rule_cross.Add( new_rule );
				RecoverReorder( n2, n );

				// reorder type4: 0 1 2  -> 1 2 0
				tmp = n2.child[0]; n2.child[0] = n2.child[1]; n2.child[1] = n2.child[2]; n2.child[2] = tmp;
				new_rule.after = "A1 A2 A0";
				new_rule.after_cross = GetNumCross( t2, n2, s );
				rule_cross.Add( new_rule );
				RecoverReorder( n2, n );

				// reorder type4: 0 1 2  -> 2 1 0
				tmp = n2.child[0]; n2.child[0] = n2.child[2]; n2.child[2] = tmp;
				new_rule.after = "A2 A1 A0";
				new_rule.after_cross = GetNumCross( t2, n2, s );
				rule_cross.Add( new_rule );
				RecoverReorder( n2, n );
			}
		}
	}
	fprintf( stderr, "\r[Extract Rule: CR   ]   %d / %d sentences processed (# Rules: %d)\n", num_sent, num_sent, rule_cross.GetNumRules() );

	rule_cross.SaveToFile( out_files.file_rule_cr.GetFP(), stree );

	return true;
}



//////////////////////////////////////////////////////////////////////////////
bool CCopycatTrain::ExtractRuleRF()
{
	for( int i=0; i<num_sent; i++ )	{
		fprintf( stderr, "\r[Extract Rule: RF   ]   %d / %d sentences processed (# Rules: %d)", i, num_sent, rule_rf.GetNumRules() );

		CSent &s = sent[i];

		// load tree
		CTree t, t2;
		t.Init(); t2.Init();
		if( t.LoadTree( s.parse ) == false )	return false;
		if( t2.LoadTree( s.parse ) == false )	return false;
		SetNodeSpan( s, t );
		SetNodeSpan( s, t2 );

		// store src statics
		for( int j=0; j<t.num_node; j++ )	{
			CNode &n = t.node[j];
			CNode &n2 = t2.node[j];

			if( MCH == "2" )	{ if( n.num_child != 2 )	continue; }
			else if( MCH == "23" )	 { if( n.num_child == 2  || n.num_child ==3 )	{ } else	{ continue; } }
			else { if( n.num_child <= 1 ) continue;  }

			if( IsAllChildrenAligned( t, n ) == false ) continue;

			CRuleRF new_rule;
			new_rule.before = GetRuleStr( t, n );
			new_rule.after =  GetRuleAfterStr( t, n );
			new_rule.freq = 1;
			rule_rf.Add( new_rule );
			RecoverReorder( n2, n );
		}

	}
	fprintf( stderr, "\r[Extract Rule: RF   ]   %d / %d sentences processed (# Rules: %d)\n", num_sent, num_sent, rule_rf.GetNumRules() );

	rule_rf.SaveToFile( out_files.file_rule_rf.GetFP(), stree, false );

	return true;
}





